

template <typename TIn, typename TKernel, typename TOut, typename TAccum>
__global__ void
convolution2D_NHWC_kernel(const TIn *input, const TKernel *kernel, TOut *output,
                          int batch, int height, int width, int in_channels,
                          int out_channels, int kh, int kw, int outHeight,
                          int outWidth, int stride_h, int stride_w,
                          int paddingPre, int paddingPost) {
  int output_y = blockIdx.y * blockDim.y + threadIdx.y; // 输出高度
  int output_x = blockIdx.x * blockDim.x + threadIdx.x; // 输出宽度
  int out_channel_idx = blockIdx.z; // 使用 blockIdx.z 来获取输出通道索引

  if (output_y < outHeight && output_x < outWidth &&
      out_channel_idx < out_channels) {
    for (int batch_idx = 0; batch_idx < batch; batch_idx++) {
      TAccum value = 0; // 使用 TAccum 类型的累加值
      for (int c = 0; c < in_channels; c++) { // 累加所有输入通道
        for (int i = 0; i < kh; i++) {
          for (int j = 0; j < kw; j++) {
            // 计算输入图像中的坐标
            int input_y = output_y * stride_h - paddingPre + i;
            int input_x = output_x * stride_w - paddingPost + j;
            if (input_y >= 0 && input_y < height && input_x >= 0 &&
                input_x < width) {
              // 获取输入值
              int input_index = batch_idx * height * width * in_channels +
                                input_y * width * in_channels +
                                input_x * in_channels + c;
              int kernel_index = out_channel_idx * kh * kw * in_channels +
                                 i * kw * in_channels + j * in_channels + c;
              value += static_cast<TAccum>(input[input_index]) *
                       static_cast<TAccum>(kernel[kernel_index]);
            }
          }
        }
      }
      // 将计算的值写入输出
      int output_index = batch_idx * outHeight * outWidth * out_channels +
                         output_y * outWidth * out_channels +
                         output_x * out_channels + out_channel_idx;
      output[output_index] = static_cast<TOut>(value); // 转换为输出类型
    }
  }
}